import vector
import matrix

from OpenGL.GL import *

class Skel:
    bone = []
    vertex = []
    def __init__(self, vertex, bone):
        """Represents a skeleton. vertex is a list of vertex data and bone is a list
        of pairs of vertex indices."""
        self.vertex = vertex
        self.bone = bone
        self.neighbor = {}
        for i in range (len(vertex)):
            self.neighbor [i] = []
        for b in bone:
            self.neighbor [b[0]].append(b[1])
            self.neighbor [b[1]].append(b[0])
    
    def __repr__(self):
        """An ascii representation of this class."""
        vrep = [str(i)+": "+str(self.vertex[i]) for i in range(len(self.vertex))]
        brep = [str(i)+": "+str(self.neighbor[i]) for i in self.neighbor]
        return "Vertices\n"+"\n".join(vrep)+"\nBones\n"+"\n".join(brep)
        
    def drawvertices (self,highlight=[]):
        """Draws vertices of the skeleton. If highlight is non-empty,
        only those vertices will be drawn. Otherwise, all vertices are 
        drawn."""
        if highlight==[]: highlight = range(len(self.vertex))
        glBegin(GL_POINTS)
        for i in highlight: glVertex3fv(self.vertex[i])
        glEnd()

    def drawbones (self):
        """Draws the bones of the skeleton."""        
        glBegin(GL_LINES)
        for x in self.neighbor:
            for y in self.neighbor[x]:
                if x<y:
                    glVertex3fv(self.vertex[x])
                    glVertex3fv(self.vertex[y])
        glEnd()
                
    def draw (self):
        """Performs an opengl rendering of the skeleton."""
        self.drawvertices()
        self.drawbones()
        
    def box (self):
        """Returns a tuple with the minimum and maximum vertex of a bounding
        box of all vertices."""
        minvtx = reduce(lambda a,b: [min(a[0],b[0]),
                                     min(a[1],b[1]),
                                     min(a[2],b[2])], self.vertex)
        maxvtx = reduce(lambda a,b: [max(a[0],b[0]),
                                     max(a[1],b[1]),
                                     max(a[2],b[2])], self.vertex)
        return minvtx,maxvtx
        
    def degree (self, i):
        """Returns the degree of node i."""
        return len(self.neighbor[i])
        
    def leaves (self):
        """Returns a list of degree 1 nodes (graph leaves)."""
        return [i for i in range(len(self.vertex)) if self.degree(i)==1]
        
    def trunks (self):
        """Returns a list of nodes of degree greater than 2."""
        return [i for i in range(len(self.vertex)) if self.degree(i)>2]
        
    def path (self, i, j, visited=[]):
        """Returns a path between vertex i and vertex j."""
        if j in self.neighbor[i]: return [i,j]
        for k in self.neighbor[i]:
            if k not in visited:
                path2 = self.path(k,j,visited+[i])
                if path2 != []: return [i]+path2
        return []

    def trunkdestinations (self, i):
        """Returns a list of truns or leaves directly connected to trunk i."""
        t = self.trunks()
        l = self.leaves()
        return [x[-1] for x in self.paths(i, set(t+l))]

    def trunkleaves (self, i):
        """Returns a list of leaves directly connected to trunk i."""
        t = self.trunks()
        l = self.leaves()
        return [x[-1] for x in self.paths(i, set(t+l)) if self.degree(x[-1])==1]
        
    def terminaltrunks (self):
        """Returns the trunk nodes which are connected to no more than one
        other trunk node."""
        def isterminal (i):
            return len([x for x in self.trunkdestinations(i) if x in t])<2
        t = self.trunks()
        return [x for x in t if isterminal(x)]
        
    def paths (self, i, goal=set([]), visited=set([])):
        """Returns all paths stemming from node i to any node in goal.
        visited is used recursively to keep track of visited nodes."""
        result = []
        for j in self.neighbor[i]:
            if j not in visited:
                if j in goal: result.append([i,j])
                else: 
                    visited.add(i)
                    for p in self.paths (j, goal, visited):
                        result.append ([i]+p)
                    visited.remove(i)
        return result

def matchskel (s1, s2): 
    """Given two skeletons, returns a list of pairs (a,b) where a and
    b are vertex indices of s1 and s2 respectively. Only trunk and 
    leaf nodes are mapped."""
    
    def extreme(s,l,v,origin=(0,0,0)):
        """Given a skeleton s, a list of nodes l and a vector v,
        returns a tuple (node,dist), where node which is furthest in the direction of v
        and dist is the distance metric (actually, a dot product)."""
        distmax = -1e100
        nodemax = -1
        for node in l:
            branchvector = vector.sub(s.vertex[node],origin)
            sz = vector.length(branchvector)
            if sz > 1.0e-100: branchvector = vector.normalize(branchvector)
            dist = vector.dot(branchvector,v)
            if dist>distmax:
                nodemax,distmax = node,dist
        return (nodemax,distmax)
        
    def aligns1tos2 ():
        """Transforms all vertices of s1 so that vector s1bottom-s1top is
        aligned with s2bottom-s2top."""
        v1 = vector.sub(s1.vertex[s1top],s1.vertex[s1bottom])
        v2 = vector.sub(s2.vertex[s2top],s2.vertex[s2bottom])
        ox,oy,oz = s1.vertex[s1bottom]
        m = matrix.product(
               matrix.translation(ox,oy,oz),
               matrix.product(
                  matrix.rotby2vectors(v1,v2),
                  matrix.translation(-ox,-oy,-oz)))
        for i in range(len(s1.vertex)):
            s1.vertex[i]=matrix.transformpoint(m,s1.vertex[i])
            
    tt1 = s1.terminaltrunks()
    tt2 = s2.terminaltrunks()
    s1top = extreme(s1,tt1,(0.0,1.0,0.0))[0]
    s1bottom = extreme(s1,tt1,(0.0,-1.0,0.0))[0]
    s2top = extreme(s2,tt2,(0.0,1.0,0.0))[0]
    s2bottom = extreme(s2,tt2,(0.0,-1.0,0.0))[0]

    # first two matches are the topmost and bottommost trunk nodes in each skeleton
    trunkmatches = [(s1top,s2top),(s1bottom,s2bottom)]
    # find the leaf node matches for each trunk
    aligns1tos2()
    leafmatches = []
    for t1,t2 in trunkmatches:
        l1 = s1.trunkleaves(t1)
        l2 = s2.trunkleaves(t2)
        origin1 = s1.vertex[t1]
        origin2 = s2.vertex[t2]
        # first, order matches by order of preference
        l1ordered = []
        for il1 in l1:
            v = vector.sub(s1.vertex[il1],origin1)
            l1ordered.append((il1,)+extreme(s2,l2,v,origin2))
        l1ordered.sort(key=lambda x:x[2],reverse=True)
        l1 = [x[0] for x in l1ordered]
        # find matches for leaves in order of preference
        for il1 in l1:
            if l2 == []: break
            v = vector.sub(s1.vertex[il1],origin1)
            il2 = extreme(s2,l2,v,origin2)[0]
            l2.remove(il2)
            leafmatches += [(il1,il2)]
    return trunkmatches+leafmatches
    
def readskel(filename):
    """Reads a file of type .ske and returns a Skel object."""
    f = open(filename)
    line = f.read().split("\n")
    assert (line[0][:3]=="SKE")
    nvtx = int(line[1])
    vtx = [[float(x) for x in row.split()] for row in line[2:2+nvtx]]
    i = 0
    bone = []
    for connect in line[2+nvtx:2+nvtx+nvtx]:
        for j in [int(x) for x in connect.split()][1:]:
            if j>i: bone.append((i,j))
        i+=1
    return Skel(vtx,bone)

